#!/usr/bin/env python3
"""
Redis MCP Server

面向 Redis 数据库的命令行操作，提供以下 MCP tools：

1. redis_benchmark_set - 执行 redis-benchmark 的 set 操作性能测试
   参数：iterations (测试迭代次数), key_range (键范围)
   返回：统一的 JSON 结构

2. redis_benchmark_multi - 执行多命令（ping,set,get）的性能测试
   参数：iterations (测试迭代次数), output_format (输出格式)
   返回：统一的 JSON 结构

3. redis_benchmark_eval - 执行 Lua 脚本的性能测试
   参数：iterations (测试迭代次数), key_range (键范围)
   返回：统一的 JSON 结构

4. redis_benchmark_lpush - 执行列表推送操作的性能测试
   参数：iterations (测试迭代次数), key_range (键范围)
   返回：统一的 JSON 结构

5. redis_cli_set_get - 设置和获取键值对
   参数：host (主机), port (端口), key (键), value (值)
   返回：统一的 JSON 结构

6. redis_cli_bulk_operations - 批量执行列表操作
   参数：repeat_count (重复次数), list_name (列表名), value (值)
   返回：统一的 JSON 结构

7. redis_cli_monitoring - 监控 Redis 内存使用情况
   参数：repeat_count (重复次数), interval (间隔秒数)
   返回：统一的 JSON 结构

8. redis_cli_scan - 扫描匹配模式的键
   参数：pattern (匹配模式)
   返回：统一的 JSON 结构

9. redis_cli_ping - 测试 Redis 连接状态
   参数：host (主机), port (端口), repeat_count (重复次数), interval (间隔秒数)
   返回：统一的 JSON 结构

10. redis_cli_interactive - 执行交互式命令
    参数：host (主机), port (端口), command (命令)
    返回：统一的 JSON 结构

11. redis_check_rdb - 检查 RDB 文件完整性
    参数：rdb_file (RDB 文件路径)
    返回：统一的 JSON 结构
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("Redis MCP Server")

@mcp.tool()
def redis_benchmark_set(iterations: int = 1000000, key_range: int = 100000000) -> dict:
    """
    执行 redis-benchmark 的 set 操作性能测试
    
    参数:
        iterations: 测试迭代次数，默认 1000000
        key_range: 键的范围大小，默认 100000000
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-benchmark -t set -n {iterations} -r {key_range}"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_benchmark_multi(iterations: int = 100000, output_format: str = "csv") -> dict:
    """
    执行多命令（ping,set,get）的性能测试
    
    参数:
        iterations: 测试迭代次数，默认 100000
        output_format: 输出格式，支持 csv 或其他格式
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-benchmark -t ping,set,get -n {iterations} --{output_format}"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_benchmark_eval(iterations: int = 10000, key_range: int = 10000) -> dict:
    """
    执行 Lua 脚本的性能测试
    
    参数:
        iterations: 测试迭代次数，默认 10000
        key_range: 键的范围大小，默认 10000
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-benchmark -r {key_range} -n {iterations} eval 'return redis.call(\"ping\")' 0"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_benchmark_lpush(iterations: int = 10000, key_range: int = 10000) -> dict:
    """
    执行列表推送操作的性能测试
    
    参数:
        iterations: 测试迭代次数，默认 10000
        key_range: 键的范围大小，默认 10000
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-benchmark -r {key_range} -n {iterations} lpush mylist __rand_int__"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_cli_set_get(host: str = "127.0.0.1", port: int = 6379, key: str = "hello", value: str = "world") -> dict:
    """
    设置和获取键值对
    
    参数:
        host: Redis 主机地址，默认 127.0.0.1
        port: Redis 端口，默认 6379
        key: 要设置的键名
        value: 要设置的值
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    # 设置键值
    set_command = f"redis-cli -h {host} -p {port} set {key} {value}"
    set_result = subprocess.run(set_command, shell=True, capture_output=True, text=True)
    
    # 获取值
    get_command = f"redis-cli -h {host} -p {port} get {key}"
    get_result = subprocess.run(get_command, shell=True, capture_output=True, text=True)
    
    return {
        "success": set_result.returncode == 0 and get_result.returncode == 0,
        "command": f"{set_command} && {get_command}",
        "exit_code": set_result.returncode if set_result.returncode != 0 else get_result.returncode,
        "stdout": f"SET: {set_result.stdout}\nGET: {get_result.stdout}",
        "stderr": f"SET: {set_result.stderr}\nGET: {get_result.stderr}"
    }

@mcp.tool()
def redis_cli_bulk_operations(repeat_count: int = 100, list_name: str = "mylist", value: str = "x") -> dict:
    """
    批量执行列表操作
    
    参数:
        repeat_count: 重复执行次数，默认 100
        list_name: 列表名称，默认 mylist
        value: 要推送的值，默认 x
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-cli -r {repeat_count} lpush {list_name} {value}"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_cli_monitoring(repeat_count: int = 100, interval: int = 1) -> dict:
    """
    监控 Redis 内存使用情况
    
    参数:
        repeat_count: 重复执行次数，默认 100
        interval: 执行间隔秒数，默认 1
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-cli -r {repeat_count} -i {interval} info | grep used_memory_human:"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_cli_scan(pattern: str = "*:12345*") -> dict:
    """
    扫描匹配模式的键
    
    参数:
        pattern: 匹配模式，默认 *:12345*
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-cli --scan --pattern '{pattern}'"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_cli_ping(host: str = "127.0.0.1", port: int = 6379, repeat_count: int = 3, interval: int = 1) -> dict:
    """
    测试 Redis 连接状态
    
    参数:
        host: Redis 主机地址，默认 127.0.0.1
        port: Redis 端口，默认 6379
        repeat_count: 重复执行次数，默认 3
        interval: 执行间隔秒数，默认 1
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    if interval > 0:
        command = f"redis-cli -h {host} -p {port} -r {repeat_count} -i {interval} ping"
    else:
        command = f"redis-cli -h {host} -p {port} -r {repeat_count} ping"
    
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_cli_interactive(host: str = "127.0.0.1", port: int = 6379, command: str = "bgsave") -> dict:
    """
    执行交互式命令
    
    参数:
        host: Redis 主机地址，默认 127.0.0.1
        port: Redis 端口，默认 6379
        command: 要执行的命令，默认 bgsave
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    expect_script = f"""
    set timeout 10
    log_file testlog
    spawn redis-cli -h {host} -p {port}
    expect {{
        "{host}:{port}>" {{ send "{command}\\\\r"
        expect "{host}:{port}>" {{send "exit\\\\r"}}
    }}
    }}
    expect eof
    """
    
    full_command = f"expect -c '{expect_script}'"
    result = subprocess.run(full_command, shell=True, capture_output=True, text=True)
    
    # 检查日志文件是否有错误
    check_command = "grep -iE \"error|failed\" testlog"
    check_result = subprocess.run(check_command, shell=True, capture_output=True, text=True)
    
    success = result.returncode == 0 and check_result.returncode == 1
    
    return {
        "success": success,
        "command": full_command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def redis_check_rdb(rdb_file: str = "/var/lib/redis/dump.rdb") -> dict:
    """
    检查 RDB 文件完整性
    
    参数:
        rdb_file: RDB 文件路径，默认 /var/lib/redis/dump.rdb
    
    返回:
        JSON 结构: {success, command, exit_code, stdout, stderr}
    """
    command = f"redis-check-rdb {rdb_file}"
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

if __name__ == "__main__":
    mcp.run()